import type { VariableStore } from "@typebot.io/forge/types";
import { isDefined, isEmpty } from "@typebot.io/lib/utils";
import type { CoreAssistantMessage, CoreMessage, CoreUserMessage } from "ai";
import { splitUserTextMessageIntoBlocks } from "./splitUserTextMessageIntoBlocks";
import type { DialogueMessage, Message, StandardMessage } from "./types";

type Props = {
  messages: Message[] | undefined;
  isVisionEnabled: boolean;
  shouldDownloadImages: boolean;
  variables: VariableStore;
};

export const parseChatCompletionMessages = async ({
  messages,
  isVisionEnabled,
  shouldDownloadImages,
  variables,
}: Props): Promise<CoreMessage[]> => {
  if (!messages) return [];
  const parsedMessages: CoreMessage[] = (
    await Promise.all(
      messages.map(async (message) => {
        if (!message.role) return;

        if (message.role === "Dialogue")
          return parseDialogueMessage({
            message,
            variables,
            isVisionEnabled,
            shouldDownloadImages,
          });

        return parseStandardMessage({
          message,
          variables,
          isVisionEnabled,
          shouldDownloadImages,
        });
      }),
    )
  )
    .flat()
    .filter(isDefined);

  return parsedMessages;
};

const parseDialogueMessage = async ({
  message,
  variables,
  isVisionEnabled,
  shouldDownloadImages,
}: Pick<Props, "variables" | "isVisionEnabled" | "shouldDownloadImages"> & {
  message: DialogueMessage;
}) => {
  if (!message.dialogueVariableId) return;
  const dialogue = variables.get(message.dialogueVariableId) ?? [];
  const dialogueArr = Array.isArray(dialogue) ? dialogue : [dialogue];

  return Promise.all(
    dialogueArr.map<
      Promise<CoreUserMessage | CoreAssistantMessage | undefined>
    >(async (dialogueItem, index) => {
      if (!dialogueItem) return;
      if (index === 0 && message.startsBy === "assistant")
        return { role: "assistant" as const, content: dialogueItem };
      if (index % (message.startsBy === "assistant" ? 1 : 2) === 0) {
        return {
          role: "user" as const,
          content: isVisionEnabled
            ? await splitUserTextMessageIntoBlocks({
                input: dialogueItem ?? "",
                shouldDownloadImages,
              })
            : dialogueItem,
        };
      }
      return { role: "assistant" as const, content: dialogueItem };
    }),
  );
};

const parseStandardMessage = async ({
  message,
  variables,
  isVisionEnabled,
  shouldDownloadImages,
}: Pick<Props, "variables" | "isVisionEnabled" | "shouldDownloadImages"> & {
  message: StandardMessage;
}) => {
  if (!message.content) return;

  const content = variables.parse(message.content);

  if (isEmpty(content)) return;

  if (message.role === "user")
    return {
      role: "user" as const,
      content: isVisionEnabled
        ? await splitUserTextMessageIntoBlocks({
            input: content,
            shouldDownloadImages,
          })
        : content,
    };

  return {
    role: message.role,
    content,
  };
};
